#  -*- coding: utf-8 -*-

from data.data_module import DataModule
import traceback
from trading.account import Account
from util.database import DB_CONN
from pathlib import Path
import pandas as pd
from util.stock_util import judge_code_trading_date,get_trading_dates,get_diff_dates,get_sub_industry,get_choice_block_mom_value,calc_negative_diff_dates,calc_positive_diff_dates
import copy,time
from pymongo import UpdateOne
from datetime import datetime
from util.database import base_code_path

"""
模拟交易数据处理，用于验证交易思路，如卖出策略等
"""

#Flag用于判断该卖出条件只使用一次
sell_strategy_1 = {
    'name' : "s1",
    'steps' : 3,
    #涨幅达到一定程度才触发卖出条件
    'init_condition': 15,
    'sell_point_condition_1' : 'upper_shadow',
    'sell_point_ratio_1' : 1/3,
    'sell_point_condition_2': 'ma_5_turn_down',
    'sell_point_ratio_2': 1/3,
    'sell_point_condition_3':'touch_ma_20',
    'sell_point_ratio_3' : 1/3,
    
}

sell_strategy_2= {
    'name': "s2",
    'steps': 3,
    # 涨幅达到一定程度才触发卖出条件
    'init_condition': 15,
    'sell_point_condition_1': 'top_2_close',
    'sell_point_ratio_1': 1/3,
    'sell_point_condition_2': 'touch_ma_5',
    'sell_point_ratio_2': 2/9,
    'sell_point_condition_3': 'touch_ma_10',
    'sell_point_ratio_3': 4/9,

}


class SimulationTrading:
    def __init__(self):

        self.dm = DataModule()
        self.cur_date = datetime.now().strftime('%Y-%m-%d')
        self.file = Path(f"{base_code_path}/trading/SimulationTrading.xlsx")

    def gen_trading_record(self,strategy):
        trading_record = dict()
        for i in range(1,strategy['steps']+1):
            price_key = f"{strategy['name']}_sell_{i}_price"
            date_key = f"{strategy['name']}_sell_{i}_date"
            trading_record[price_key] = 0.0
            trading_record[date_key] = ''

        trading_record[f"{strategy['name']}_profit_ratio"] = 0.0
        trading_record[f"stop_loss_price"] = 0.0
        trading_record[f"stop_loss_date"] = ''

        return trading_record

    def touch_ma_5(self,date_df,date):
        if date_df.loc[date]['close'] <= date_df.loc[date]['ma5']:
            return date_df.loc[date]['close'],True
        else:
            return 0,False

    def touch_ma_10(self,date_df,date):
        if date_df.loc[date]['close'] <= date_df.loc[date]['ma10']:
            return date_df.loc[date]['close'],True
        else:
            return 0,False

    def touch_ma_20(self,date_df,date):
        if date_df.loc[date]['close'] <= date_df.loc[date]['ma20']:
            return date_df.loc[date]['close'],True
        else:
            return 0,False

    def upper_shadow(self,date_df,date):
        #print(date,"upper_shadow")
        close = date_df.loc[date]['close']
        high = date_df.loc[date]['high']
        pre_close = date_df.loc[date]['pre_close']
        #上影线4个点
        if round((high - close)/pre_close*100,2) > 4:
            return date_df.loc[date]['close'],True
        else:
            return 0,False

    def ma_5_turn_down(self,date_df,date):
        last_date = calc_negative_diff_dates(date_df.loc[date]['code'],False,date,-1)
        ma5  =  date_df.loc[date]['ma5']
        last_ma5  =  date_df.loc[last_date]['ma5']

        if last_ma5 > ma5:
            return date_df.loc[date]['close'],True
        else:
            return 0,False

    def gen_strategy_condition_flag(self,strategy):
        #用于判断每个卖出条件仅使用一次
        strategy_condition_flag = dict()
        strategy_condition_flag[f"{strategy['name']}_flag"] = 0
        return strategy_condition_flag


    def judge_sell_condition(self,date_df,date,i,strategy,strategy_condition_flag):
        op_dict = {
            'upper_shadow':self.upper_shadow,
            'ma_5_turn_down':self.ma_5_turn_down,
            'touch_ma_20':self.touch_ma_20,
            'top_2_close':self.upper_shadow,
            'touch_ma_5':self.touch_ma_5,
            'touch_ma_10':self.touch_ma_10,

        }
        #卖出条件按顺序触发
        if i - strategy_condition_flag[f"{strategy['name']}_flag"] != 1:
            return  0,False

        sell_price,result = op_dict.get(strategy[f'sell_point_condition_{i}'], 'wrong para')(date_df,date)
        #print(date,strategy['name'],strategy[f'sell_point_condition_{i}'],strategy_condition_flag,sell_price,result)
        if result:
            strategy_condition_flag[f"{strategy['name']}_flag"] += 1
            #print(date,sell_price,strategy[f'sell_point_condition_{i}'],result)
        return sell_price,result

    def judge_stop_loss(self,date_df,date,buy_price,strategy):
        close = date_df.loc[date]['close']
        cur_ratio = round((close - buy_price)/buy_price*100,2)

        sell_price,result = self.touch_ma_20(date_df,date)

        if cur_ratio < strategy['init_condition']:
            if result:
                #达到止损条件
                return sell_price,True
            else:
                #小于涨幅位置，但无需止损
                return 0,True
        else:
            #大于涨幅位置，后续不用计算止损了
            return 0,False

    def simulation_trading_one_stock(self,row,strategy):
        buy_date = row['buy_date']
        end_date = self.cur_date
        buy_price = row['buy_price']
        code = row['code']
        trading_record = self.gen_trading_record(strategy)
        begin_date = calc_negative_diff_dates(code,False,buy_date,-30)
        simulation_date = calc_negative_diff_dates(code,False,buy_date,-5)
        trade_date = calc_positive_diff_dates(code,False,buy_date,1)

        total_volume = 900
        cur_volume = total_volume
        init_value = buy_price * total_volume
        sell_value = 0.0
        stop_loss_date = self.cur_date
        #print(row)
        try:
            df_daily = self.dm.get_k_data(code,autype=None,begin_date=begin_date,end_date=end_date)
            if df_daily.index.size > 0:
                df_daily.set_index(['date'], 1, inplace=True)
            df_trading_daily = df_daily.loc[df_daily.is_trading == True, :]
            df_trading_daily_copy = df_trading_daily.copy()

            df_trading_daily_copy['ma5'] = round(df_trading_daily_copy['close'].rolling(5).mean(), 2)
            df_trading_daily_copy['ma10'] = round(df_trading_daily_copy['close'].rolling(10).mean(), 2)
            df_trading_daily_copy['ma20'] = round(df_trading_daily_copy['close'].rolling(20).mean(), 2)
            
            df_target_daily = df_trading_daily_copy.loc[simulation_date:end_date]
            #print(df_target_daily)
            strategy_condition_flag = self.gen_strategy_condition_flag(strategy)
            for date in df_target_daily.index:
                if date < trade_date:
                    continue

                #止损的判断逻辑，小于涨幅时，判断止损，符合则止损，不符合继续下一次，大于涨幅时继续按卖出条件计算
                if date < stop_loss_date:
                    stop_loss_price,result = self.judge_stop_loss(df_target_daily,date,buy_price,strategy)
                    if result:
                        if sell_value > 0:
                            sell_value += stop_loss_price * cur_volume
                            cur_volume = 0
                            trading_record[f"stop_loss_price"] = stop_loss_price
                            trading_record[f"stop_loss_date"] = date
                            profit_ratio = round((sell_value - init_value) / init_value * 100, 2)
                            trading_record[f"{strategy['name']}_profit_ratio"] = profit_ratio
                        else:
                            continue
                    else:
                        stop_loss_date = date

                if cur_volume == 0:
                    break

                for i in range(1,strategy['steps']+1):

                    sell_price,result = self.judge_sell_condition(df_target_daily,date,i,strategy,strategy_condition_flag)
                    if result:
                        trading_record[f"{strategy['name']}_sell_{i}_price"] = sell_price
                        trading_record[f"{strategy['name']}_sell_{i}_date"] = date
                        sell_value += sell_price * total_volume * strategy[f"sell_point_ratio_{i}"]
                        cur_volume -= total_volume * strategy[f"sell_point_ratio_{i}"]
                        #print(f"cur_volume:{cur_volume}")
                        if cur_volume == 0:
                            break
            if cur_volume == 0:
                profit_ratio = round((sell_value - init_value)/init_value *100,2)
                trading_record[f"{strategy['name']}_profit_ratio"] = profit_ratio

        except Exception as e:
            print(f"simulation_trading_one_stock error :{e}")
        
        return trading_record

    
    def sumulation_trading_all_stocks(self):
        data_df = pd.read_excel(self.file, encoding="gb2312",dtype=object)
        data_df['buy_date'] = data_df['buy_date'].astype(str)
        data_df['code'] = data_df['code'].astype(str)
        data_df['stop_loss_price'] = 0.0
        data_df['stop_loss_date'] = ""
        #print(data_df)
        for strategy in [sell_strategy_1,sell_strategy_2]:
            for i in range(1,strategy['steps']+1):
                data_df[f"{strategy['name']}_sell{i}_price"] = 0.0
                data_df[f"{strategy['name']}_sell{i}_date"] = ''
            data_df[f"{strategy['name']}_profit_ratio"] = 0.0


            for index, row in data_df.iterrows():
                trading_record = self.simulation_trading_one_stock(row,strategy)
                #print(trading_record)
                for i in range(1,strategy['steps']+1):
                    data_df.loc[index,f"{strategy['name']}_sell{i}_price"] = trading_record[f"{strategy['name']}_sell_{i}_price"]
                    data_df.loc[index,f"{strategy['name']}_sell{i}_date"] = trading_record[f"{strategy['name']}_sell_{i}_date"]
                data_df.loc[index,f"{strategy['name']}_profit_ratio"] = trading_record[f"{strategy['name']}_profit_ratio"]
                data_df['stop_loss_price'] = trading_record['stop_loss_price']
                data_df['stop_loss_date'] = trading_record['stop_loss_date']

        data_df.to_csv("simulation_trading.csv")

if __name__ == '__main__':
    pd.set_option('display.width',500)
    pd.set_option('display.max_columns', 500)
    pd.set_option('display.max_colwidth', 500)
    st = SimulationTrading()
    start_time = time.time()
    st.sumulation_trading_all_stocks()
    end_time = time.time()
    print(f"用时:{end_time - start_time}")